{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import hail as hl \n",
    "hl.init(log='./logs/20221025.log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from statsmodels import robust\n",
    "from statistics import median"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Files\n",
    "Hail需要输入.bgz格式的vcf文件，可以使用bgzip和tabix生成"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vcf_file = ''\n",
    "working_dir = ''\n",
    "prefix = ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ref_dir = '../ref/hail/'\n",
    "\n",
    "mt_raw_file = working_dir + '/data/' + prefix + '.raw.mt'\n",
    "mt_splitmulti_file = working_dir + '/data/' + prefix + '.split.mt'\n",
    "mt_sampleqc_file = working_dir + '/data/' + prefix + '.qc.mt'\n",
    "mt_variantqc_file= working_dir + '/data/'+ prefix + '.vq.mt'\n",
    "\n",
    "sample_qc_info_file = working_dir + '/info/' + prefix + '_sample_qc_info.csv'\n",
    "variant_qc_info_file = working_dir + '/info/' + prefix + '_variant_qc_info.csv'\n",
    "\n",
    "pca_value_file = working_dir + '/data/'+ prefix + '.pcaval.txt'\n",
    "pca_score_file = working_dir + '/data/'+ prefix + '.pcavec.txt'\n",
    "\n",
    "vcf_qc_file = working_dir + '/qc/'+ prefix + '.HC.VQSR.qc.vcf'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Reference Genome "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ref = hl.ReferenceGenome.read(ref_dir + 'ensembl_GRCh38.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### vcf -> mt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = hl.import_vcf(vcf_file, reference_genome=ref).write(mt_raw_file, overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = hl.read_matrix_table(mt_raw_file)\n",
    "count_sample_raw_0 = mt.count_cols()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Split Multi-allelic sites into biallelic sites\n",
    "https://hail.is/docs/0.2/methods/genetics.html#hail.methods.split_multi_hts \\\n",
    "https://gatk.broadinstitute.org/hc/en-us/articles/360035890771-Biallelic-vs-Multiallelic-sites"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = hl.split_multi_hts(mt.select_entries(mt.GT, mt.AD, mt.DP, mt.GQ, mt.PL)).naive_coalesce(20000).write(mt_splitmulti_file, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = hl.read_matrix_table(mt_splitmulti_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1 Sample QC"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.1 Outlier Detection\n",
    "**根据样本量大小适当调整阈值，小样本量群体建议跳过此步骤** \\\n",
    "https://hail.is/docs/0.2/methods/genetics.html#hail.methods.sample_qc \\\n",
    "VCF-based sample-level QC metrics used for outlier detection and computed with GATK included: call rate, number of SNPs, number of deletions, number of insertions, insertion to deletion ratio, transition to transversion ratio, heterozygous to homozygous ratio. All samples had a call rate above 99%. Samples that were 4 median absolute deviations above or below the median of any of the above QC metrics were manually inspected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = hl.sample_qc(mt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#样本个数、样本名称、QC指标名称\n",
    "nums = mt.count_cols()\n",
    "samples = list(mt.s.take(nums))\n",
    "metrics = ['call_rate', 'n_snp', 'n_deletion', 'n_insertion', 'r_insertion_deletion', 'r_ti_tv', 'r_het_hom']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#从mt中提取各样本的QC指标\n",
    "call_rate = list(mt.sample_qc.call_rate.take(nums))\n",
    "n_snp = list(mt.sample_qc.n_snp.take(nums))\n",
    "n_deletion = list(mt.sample_qc.n_deletion.take(nums))\n",
    "n_insertion = list(mt.sample_qc.n_insertion.take(nums))\n",
    "r_insertion_deletion = list(mt.sample_qc.r_insertion_deletion.take(nums))\n",
    "r_ti_tv = list(mt.sample_qc.r_ti_tv.take(nums))\n",
    "r_het_hom = list(mt.sample_qc.r_het_hom_var.take(nums))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_qc = pd.DataFrame([call_rate, n_snp, n_deletion, n_insertion, r_insertion_deletion, r_ti_tv, r_het_hom], index=metrics, columns=samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_qc.to_csv(sample_qc_info_file, index=True, header=True, sep='\\t')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_samples = []\n",
    "for i in sample_qc.index:\n",
    "    MAD = robust.mad(sample_qc.loc[i])\n",
    "    med = median(sample_qc.loc[i])\n",
    "    for j in sample_qc.columns:\n",
    "        value = sample_qc.loc[i,j]\n",
    "        if (value > (med + 6*MAD) or value < (med - 6*MAD)):\n",
    "            filtered_samples.append(j)\n",
    "            print(i + ':' + j)\n",
    "            \n",
    "filtered_samples = list(set(filtered_samples))\n",
    "print(filtered_samples)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.2 LD Pruning\n",
    "https://gwaslab.com/2021/05/18/pruning-and-clumping/ \\\n",
    "https://hail.is/docs/0.2/methods/genetics.html#hail.methods.ld_prune"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pruned_variant_table = hl.ld_prune(mt.GT, r2=0.2, bp_window_size=500000)\n",
    "mt_pruned = mt.filter_rows(hl.is_defined(pruned_variant_table[mt.row_key]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.3 Sex-check\n",
    "https://gwaslab.com/2021/06/27/sex-check/\\\n",
    "https://hail.is/docs/0.2/methods/genetics.html#hail.methods.impute_sex\\\n",
    "性别错误 (Sex discrepancy)：有时数据录入会有错误，我们应基于X染色体的杂合性，确定是否有性别录入错误的个体。PLINK默认对于男性，X染色体的纯合性（homozygosity）估计值应高于0.8，女性应低于0.2。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imputed_sex = hl.impute_sex(mt_pruned.GT)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.4 IBD\n",
    "https://gwaslab.com/2020/04/11/example-post-3/ \\\n",
    "https://hail.is/docs/0.2/methods/relatedness.html#hail.methods.identity_by_descent \\\n",
    "\n",
    "血缘同源（Identity By Descent，IBD），子代中共有的等位基因来源于同一祖先。\\\n",
    "IBD可以让我们了解两个体间的亲缘关系，虽然无法直接测得，但可以根据IBS以及等位基因频率的分布来推定。\\\n",
    "PLINK中使用 PI_HAT 值来推定IBD的值。该方法基于隐马尔科夫模型 hidden Markov model (HMM)，通过矩估计（method-of-moments）来计算 IBD=1, 2或0 的概率。\\\n",
    "PI_HAT：为IBD比例 , 即 P(IBD=2) + 0.5*P(IBD=1)，PI_HAT的值与对应关系如下所示：\n",
    "\n",
    "PI_HAT＝0 无亲缘关系\\\n",
    "PI_HAT＝0.25 表兄弟\\\n",
    "PI_HAT＝0.5 亲子或兄弟姐妹\\\n",
    "PI_HAT＝1 本人或同卵双胞胎\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ibd = hl.identity_by_descent(mt_pruned, maf=mt_pruned['panel_maf'], min=0.2, max=0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.5 PCA\n",
    "https://gwaslab.com/2021/03/28/pca-2/\\\n",
    "https://hail.is/docs/0.2/methods/genetics.html#hail.methods.hwe_normalized_pca\\\n",
    "主成分分析（Principal Components Analysis, PCA）是一种常用的数据降维方法，在群体遗传学中被广泛用于识别并调整样本的群体分层问题。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eigenvalues, scores, loadings = hl.hwe_normalized_pca(mt_pruned.GT, k=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with hl.utils.hadoop_open(pca_value_file, 'w') as f:\n",
    "\tfor val in eigenvalues:\n",
    "\t\tf.write(str(val) + '\\n')\n",
    "scores.flatten().export(pca_score_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2 Variant QC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = hl.read_matrix_table(mt_splitmulti_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = hl.variant_qc(mt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "count_snps_sq_0 = mt.count_rows()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.1 call rate < 0.95\n",
    "对**某个snp位点**被成功检测到的样本与所有样本比值"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = mt.filter_rows(mt.variant_qc.call_rate < 0.95, keep=False)\n",
    "count_snps_callrate_1 = mt.count_rows()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.2 lay in LCR\n",
    "Low complexity regions (LCRs) consist of locally repetitive sections of the genome. Heng's paper identified these using mdust and provides a BED file of LCRs covering 2% of the genome. Repeats in these regions can lead to artifacts in sequencing and variant calling. Heng's paper provides examples of areas where a full de-novo assembly correctly resolves the underlying structure, while local reassembly variant callers do not. \\\n",
    "http://bcb.io/2014/05/12/wgs-trio-variant-evaluation/#:~:text=For%20SNPs%2C%20removing%20low%20complexity%20regions%20removes%20approximately,to%20the%20over-representation%20of%20indels%20in%20repeat%20regions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lcr_file = ref_dir + 'LCR.bed'\n",
    "lcr = hl.import_locus_intervals(lcr_file, reference_genome=ref)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = mt.filter_rows(hl.is_defined(lcr[mt.locus]), keep=False)\n",
    "count_snps_lcr_2 = mt.count_rows()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.3 Inbreeding Cofficient\n",
    "https://gatk.broadinstitute.org/hc/en-us/articles/360035531992-Inbreeding-Coefficient\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = mt.annotate_rows(F = 1.0 - (mt.variant_qc.n_het / (2.0 * (1 - mt.variant_qc.AF[1]) * mt.variant_qc.AF[1] * mt.count_cols() ) ) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = mt.filter_rows(mt.F < -0.3, keep=False)\n",
    "count_snps_F_3 = mt.count_rows()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.4 hardy-weinberg test\n",
    "- **条件**\n",
    "  - 无限大的群体\n",
    "  - 随机交配\n",
    "  - 没有突变\n",
    "  - 没有选择\n",
    "  - 没有迁移\n",
    "  - 没有遗传漂变\n",
    "- **结论**\n",
    "  - 群体内一个位点上的基因型频率和基因频率在遗传过程中保持不变，处于遗传平衡状态\n",
    "- **公式**\n",
    "  - p2+2pq+q2=1,p+q=1,(p+q)2=1\n",
    "  \n",
    "GWAS假设样本群体是符合哈温平衡的， 对于不符合哈温平衡的SNP位点，需要过滤掉\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = mt.filter_rows(mt.variant_qc.p_value_hwe < 0.000001, keep=False)\n",
    "count_snps_HWE_4 = mt.count_rows()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.5 Statistic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "snp_indexs = ['Raw', 'Call Rate < 0.95', 'Lay in Low Complexity Regions', 'InbreedingCoeff < -0.3', 'HWE < 10^(-6)']\n",
    "snp_cols = ['# PASS']\n",
    "snp_counts = [count_snps_sq_0, count_snps_callrate_1, count_snps_lcr_2, count_snps_F_3, count_snps_HWE_4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "variant_qc = pd.DataFrame(snp_counts, index=snp_indexs, columns=snp_cols)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "variant_qc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "variant_qc.to_csv(variant_qc_info_file, index=True, header=True, sep='\\t')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3 Genotype QC\n",
    "- DP < 10\n",
    "- DP > 400\n",
    "- GQ < 25\n",
    "- HET AB < 0.25\n",
    "- HOM REF AB > 0.1\n",
    "- HOM ALT AB < 0.9 \n",
    "\n",
    "AB(Allele Balance) \\\n",
    "Next is genotype QC. It’s a good idea to filter out genotypes where the reads aren’t where they should be: if we find a genotype called homozygous reference with >10% alternate reads, a genotype called homozygous alternate with >10% reference reads, or a genotype called heterozygote without a ref / alt balance near 1:1, it is likely to be an error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "count_geno_raw_0 = mt.count_cols() * mt.count_rows()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ab = mt.AD[1]/hl.sum(mt.AD)\n",
    "\n",
    "filter_condition_ab = (\n",
    "    (mt.DP < 10) |\n",
    "    (mt.DP > 400) |\n",
    "    (mt.GQ < 25) |\n",
    "    (mt.GT.is_het() & (ab < 0.25)) |\n",
    "    (mt.GT.is_het() & (ab > 0.75)) |\n",
    "    (mt.GT.is_hom_ref() & (ab > 0.1)) |\n",
    "    (mt.GT.is_hom_var() & (ab < 0.9))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fraction_filtered = mt.aggregate_entries(hl.agg.fraction(filter_condition_ab))\n",
    "print(f'Filtering {fraction_filtered * 100:.2f}% entries out of downstream analysis.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mt = mt.filter_entries(filter_condition_ab, keep=False)\n",
    "count_geno_gq_1 = count_geno_raw_0 * (1.0 - fraction_filtered)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(count_geno_raw_0, count_geno_gq_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4 Export VCF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hl.export_vcf(mt,'/home/lilabguest7/lilab5_cxf/wgs_yanke/qc/wgs_yanke.HC.VQSR.qc.vcf')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.12 ('base': conda)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "5d9cbfe4c2beda62a166aa7e4823b29b05988cb4f38d93c470931396b608585d"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
